
import sys
import os
from typing import List, Tuple

from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QTextEdit, QFileDialog, QLabel, QGridLayout,
    QDoubleSpinBox
)
from PyQt6.QtCore import Qt
import pyqtgraph as pg
import numpy as np
from obspy.core.utcdatetime import UTCDateTime
from datetime import datetime, timezone

from eps2_to_sac import read_eps2_data, create_sac_trace

pg.setConfigOption('background', 'w')
pg.setConfigOption('foreground', 'k')

class NoWheelViewBox(pg.ViewBox):
    def wheelEvent(self, ev):
        ev.accept()
        return

class EPS2ViewerPyQtGraph(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("EPS2 三分量查看器 + 截取 + 导出SAC (PyQtGraph)")
        self.setGeometry(80, 80, 1400, 820)

        # 数据与元信息
        self.metadata = None
        self.data_arrays: List[np.ndarray] = []  # 当前数据 [Z, X, Y]
        self.original_arrays: List[np.ndarray] = []  # 原始数据
        self.delta: float = 0.01  # 默认 100 sps
        self.start_utctime: UTCDateTime | None = None
        self.original_start_utctime: UTCDateTime | None = None
        # 本地时间戳基准（秒 since epoch，用于 DateAxisItem 显示本地时间）
        self.start_local_ts: float | None = None

        # 视图组件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        grid = QGridLayout(central_widget)

        # 顶部控制条
        controls = QWidget()
        ctl = QHBoxLayout(controls)
        self.load_btn = QPushButton("加载EPS2文件(.012/.013)")
        self.apply_cut_btn = QPushButton("应用截取")
        self.export_btn = QPushButton("导出当前数据为SAC(三分量)")
        self.reset_data_btn = QPushButton("重置原始数据")
        self.file_label = QLabel("未加载文件")
        self.file_label.setStyleSheet("color: gray;")

        # 轴向缩放控件
        self.zoom_in_x_btn = QPushButton("时间轴放大")
        self.zoom_out_x_btn = QPushButton("时间轴缩小")
        self.reset_x_btn = QPushButton("时间轴重置")
        self.zoom_in_y_btn = QPushButton("数值轴放大")
        self.zoom_out_y_btn = QPushButton("数值轴缩小")
        self.reset_y_btn = QPushButton("数值轴重置")

        # 截取时间控件（秒）
        self.cut_start_spin = QDoubleSpinBox()
        self.cut_start_spin.setDecimals(3)
        self.cut_start_spin.setMinimum(0.0)
        self.cut_start_spin.setMaximum(1e9)
        self.cut_start_spin.setSingleStep(0.1)
        self.cut_start_spin.setPrefix("起始秒:")
        self.cut_duration_spin = QDoubleSpinBox()
        self.cut_duration_spin.setDecimals(3)
        self.cut_duration_spin.setMinimum(0.0)
        self.cut_duration_spin.setMaximum(1e9)
        self.cut_duration_spin.setSingleStep(0.5)
        self.cut_duration_spin.setPrefix("持续秒:")

        ctl.addWidget(self.load_btn)
        ctl.addWidget(self.apply_cut_btn)
        ctl.addWidget(self.export_btn)
        ctl.addWidget(self.reset_data_btn)
        ctl.addSpacing(12)
        ctl.addWidget(self.zoom_in_x_btn)
        ctl.addWidget(self.zoom_out_x_btn)
        ctl.addWidget(self.reset_x_btn)
        ctl.addSpacing(8)
        ctl.addWidget(self.zoom_in_y_btn)
        ctl.addWidget(self.zoom_out_y_btn)
        ctl.addWidget(self.reset_y_btn)
        ctl.addSpacing(12)
        ctl.addWidget(self.cut_start_spin)
        ctl.addWidget(self.cut_duration_spin)
        ctl.addStretch(1)

        # 三分量绘图（共用时间轴）
        axis_items = {'bottom': pg.DateAxisItem()}
        self.plotZ = pg.PlotWidget(viewBox=NoWheelViewBox(), axisItems=axis_items)
        self.plotX = pg.PlotWidget(viewBox=NoWheelViewBox())
        self.plotY = pg.PlotWidget(viewBox=NoWheelViewBox())
        for pw in (self.plotZ, self.plotX, self.plotY):
            pw.getPlotItem().showGrid(x=True, y=True, alpha=0.5)
            pw.setLabel('left', "振幅")
        self.plotZ.setLabel('bottom', "时间")
        self.plotX.setXLink(self.plotZ)
        self.plotY.setXLink(self.plotZ)

        # 在主时间轴上添加可拖拽的选择区域（X 轴时间窗）
        self.region = pg.LinearRegionItem()
        self.region.setZValue(10)
        self.plotZ.addItem(self.region)
        self.region.sigRegionChanged.connect(self.on_region_changed)

        # 头信息
        self.header_text = QTextEdit()
        self.header_text.setReadOnly(True)
        self.header_text.setLineWrapMode(QTextEdit.LineWrapMode.NoWrap)
        header_label = QLabel("EPS2/SAC 头信息:")

        grid.addWidget(controls, 0, 0, 1, 2)
        grid.addWidget(self.file_label, 1, 0, 1, 2)
        grid.addWidget(self.plotZ, 2, 0)
        grid.addWidget(self.plotX, 3, 0)
        grid.addWidget(self.plotY, 4, 0)

        header_layout = QVBoxLayout()
        header_layout.addWidget(header_label)
        header_layout.addWidget(self.header_text)
        grid.addLayout(header_layout, 2, 1, 3, 1)

        grid.setRowStretch(2, 3)
        grid.setRowStretch(3, 3)
        grid.setRowStretch(4, 3)
        grid.setColumnStretch(0, 4)
        grid.setColumnStretch(1, 2)

        # 信号连接
        self.load_btn.clicked.connect(self.load_eps2_file)
        self.apply_cut_btn.clicked.connect(self.apply_cut_window)
        self.export_btn.clicked.connect(self.export_current_sac)
        self.reset_data_btn.clicked.connect(self.reset_to_original)
        self.zoom_in_x_btn.clicked.connect(lambda: self.zoom_x(1.5))
        self.zoom_out_x_btn.clicked.connect(lambda: self.zoom_x(1/1.5))
        self.reset_x_btn.clicked.connect(self.reset_x)
        self.zoom_in_y_btn.clicked.connect(lambda: self.zoom_y(1.5))
        self.zoom_out_y_btn.clicked.connect(lambda: self.zoom_y(1/1.5))
        self.reset_y_btn.clicked.connect(self.reset_y)
        self.cut_start_spin.valueChanged.connect(self.on_spin_changed)
        self.cut_duration_spin.valueChanged.connect(self.on_spin_changed)

        # 曲线引用
        self.curveZ = None
        self.curveX = None
        self.curveY = None

    def load_eps2_file(self):
        filepath, _ = QFileDialog.getOpenFileName(
            self,
            "选择EPS2文件",
            "",
            "EPS2 Files (*.012 *.013 *.00F);;All Files (*)"
        )
        if not filepath:
            return
        try:
            data, metadata = read_eps2_data(filepath)
            self.metadata = metadata
            self.delta = 1.0 / metadata.get('sampling_rate', 100)
            # 计算起始UTC时间
            if 'datetime' in metadata:
                self.start_utctime = UTCDateTime(metadata['datetime'])
            else:
                self.start_utctime = UTCDateTime()
            self.original_start_utctime = self.start_utctime
            # 计算本地时间戳基准（用于本地时间显示）
            self.start_local_ts = datetime.fromtimestamp(float(self.start_utctime.timestamp), tz=timezone.utc).astimezone().timestamp()

            # 拆分为三分量（Z, X, Y）
            z_data = data[:, 0]
            x_data = data[:, 1]
            y_data = data[:, 2]
            self.data_arrays = [z_data, x_data, y_data]
            self.original_arrays = [z_data.copy(), x_data.copy(), y_data.copy()]

            # 设定默认截取控件范围（持续秒）
            total_seconds = len(z_data) * self.delta
            self.cut_start_spin.setMaximum(max(total_seconds - self.delta, 0.0))
            self.cut_duration_spin.setMaximum(total_seconds)
            # 初始选择区域：前 10 秒或总时长
            self.cut_start_spin.setValue(0.0)
            self.cut_duration_spin.setValue(min(10.0, total_seconds))

            self.file_label.setText(filepath)
            self.update_plots()
            self.update_header()
            self.reset_all()
            # 初始化 region 到当前 spin 设置
            self.sync_region_to_spins()
        except Exception as e:
            self.file_label.setText(f"错误: {e}")

    def build_timestamps(self, npts: int) -> np.ndarray:
        base = float(self.start_utctime.timestamp)
        return base + np.arange(npts) * self.delta

    def update_plots(self):
        if not self.data_arrays:
            return
        z, x, y = self.data_arrays
        ts = self.build_timestamps(len(z))
        # 重绘
        self.plotZ.clear(); self.plotX.clear(); self.plotY.clear()
        self.curveZ = self.plotZ.plot(ts, z, pen=pg.mkPen('k', width=1), antialias=True)
        self.curveX = self.plotX.plot(ts, x, pen=pg.mkPen('b', width=1), antialias=True)
        self.curveY = self.plotY.plot(ts, y, pen=pg.mkPen('r', width=1), antialias=True)
        # 重新添加选择区域
        self.plotZ.addItem(self.region)

    def update_header(self):
        if self.metadata is None:
            return
        info = []
        for k, v in self.metadata.items():
            info.append(f"{k}: {v}")
        info.append(f"delta: {self.delta}")
        info.append(f"start_utctime: {self.start_utctime}")
        info.append(f"npts: {len(self.data_arrays[0]) if self.data_arrays else 0}")
        self.header_text.setText("\n".join(info))

    # 视图控制
    def viewbox(self) -> pg.ViewBox:
        return self.plotZ.getViewBox()

    def reset_all(self):
        self.plotZ.enableAutoRange(axis=pg.ViewBox.XYAxes, enable=True)
        self.plotZ.autoRange()
        for pw in (self.plotX, self.plotY):
            pw.enableAutoRange(axis=pg.ViewBox.YAxis, enable=True)
            pw.updateAutoRange()

    def reset_x(self):
        vb = self.viewbox()
        vb.enableAutoRange(axis=pg.ViewBox.XAxis, enable=True)
        vb.updateAutoRange()

    def reset_y(self):
        for pw in (self.plotZ, self.plotX, self.plotY):
            vb = pw.getViewBox()
            vb.enableAutoRange(axis=pg.ViewBox.YAxis, enable=True)
            vb.updateAutoRange()

    def zoom_x(self, factor: float):
        vb = self.viewbox()
        (x0, x1), _ = vb.viewRange()
        xc = 0.5 * (x0 + x1)
        span = (x1 - x0)
        new_span = span / factor if factor > 0 else span
        vb.setXRange(xc - new_span/2, xc + new_span/2, padding=0)

    def zoom_y(self, factor: float):
        for pw in (self.plotZ, self.plotX, self.plotY):
            vb = pw.getViewBox()
            _, (y0, y1) = vb.viewRange()
            yc = 0.5 * (y0 + y1)
            span = (y1 - y0)
            new_span = span / factor if factor > 0 else span
            vb.setYRange(yc - new_span/2, yc + new_span/2, padding=0)

    # 选择区域与截取逻辑
    def on_region_changed(self):
        if not self.data_arrays or self.start_utctime is None:
            return
        x_min, x_max = self.region.getRegion()
        # 将选择的时间戳转换为相对起始秒（基于本地时间戳基准）
        base_ts = float(self.start_local_ts) if self.start_local_ts is not None else float(self.start_utctime.timestamp)
        start_sec = max(0.0, x_min - base_ts)
        end_sec = max(0.0, x_max - base_ts)
        duration_sec = max(0.0, end_sec - start_sec)
        # 更新 SpinBox，不触发递归；临时断开信号
        self.cut_start_spin.blockSignals(True)
        self.cut_duration_spin.blockSignals(True)
        self.cut_start_spin.setValue(start_sec)
        self.cut_duration_spin.setValue(duration_sec)
        self.cut_start_spin.blockSignals(False)
        self.cut_duration_spin.blockSignals(False)

    def on_spin_changed(self):
        # 将 SpinBox 的值同步到图上的选择区域
        if self.start_utctime is None:
            return
        start_sec = float(self.cut_start_spin.value())
        duration_sec = float(self.cut_duration_spin.value())
        x0 = float(self.start_utctime.timestamp) + start_sec
        x1 = x0 + max(0.0, duration_sec)
        self.region.setRegion((x0, x1))

    def sync_region_to_spins(self):
        # 初始化 region 为当前 spin 的窗口
        self.on_spin_changed()

    def apply_cut_window(self):
        # 根据当前选择区域/SpinBox，截取数据并更新开始时间与样点数
        if not self.data_arrays or self.metadata is None or self.start_utctime is None:
            self.file_label.setText("错误: 请先加载EPS2文件")
            return
        start_sec = float(self.cut_start_spin.value())
        duration_sec = float(self.cut_duration_spin.value())
        if duration_sec <= 0:
            self.file_label.setText("错误: 持续秒必须大于0")
            return
        start_idx = int(round(start_sec / self.delta))
        end_idx = start_idx + int(round(duration_sec / self.delta))
        z_len = len(self.data_arrays[0])
        start_idx = max(0, min(start_idx, z_len - 1))
        end_idx = max(start_idx + 1, min(end_idx, z_len))
        # 截取数据
        self.data_arrays = [arr[start_idx:end_idx] for arr in self.data_arrays]
        # 更新起始时间（SAC头将使用该窗口起点）
        self.start_utctime = self.original_start_utctime + start_idx * self.delta
        # 同步本地时间戳基准
        self.start_local_ts = datetime.fromtimestamp(float(self.start_utctime.timestamp), tz=timezone.utc).astimezone().timestamp()
        # 刷新视图与头信息
        self.update_plots()
        self.update_header()
        self.file_label.setText("已应用截取窗口，当前数据为所选时间段")

    def reset_to_original(self):
        # 恢复原始数据与起始时间
        if not self.original_arrays:
            return
        self.data_arrays = [arr.copy() for arr in self.original_arrays]
        self.start_utctime = self.original_start_utctime
        self.update_plots()
        self.update_header()
        self.file_label.setText("已恢复至原始数据")
        # 重置 region 与 spin
        total_seconds = len(self.data_arrays[0]) * self.delta
        self.cut_start_spin.setMaximum(max(total_seconds - self.delta, 0.0))
        self.cut_duration_spin.setMaximum(total_seconds)
        self.cut_start_spin.setValue(0.0)
        self.cut_duration_spin.setValue(min(10.0, total_seconds))
        self.sync_region_to_spins()

    # 导出当前数据为 SAC
    def export_current_sac(self):
        if not self.data_arrays or self.metadata is None or self.start_utctime is None:
            self.file_label.setText("错误: 当前无数据可导出")
            return
        out_dir = QFileDialog.getExistingDirectory(self, "选择SAC输出目录")
        if not out_dir:
            return
        # 设置 metadata 开始时间为当前数据起点
        md = dict(self.metadata)
        md['datetime'] = self.start_utctime.datetime
        components = ['Z', 'X', 'Y']
        sac_paths: List[str] = []
        for comp, arr in zip(components, self.data_arrays):
            tr = create_sac_trace(arr.astype(np.float32), md, component=comp)
            base_name = f"{int(self.start_utctime.timestamp)}_{comp}.sac"
            out_path = os.path.join(out_dir, base_name)
            tr.write(out_path, format='SAC')
            sac_paths.append(out_path)
        self.file_label.setText(f"已导出 {len(sac_paths)} 个SAC:\n" + "\n".join(sac_paths))


def main():
    app = QApplication(sys.argv)
    viewer = EPS2ViewerPyQtGraph()
    viewer.show()
    sys.exit(app.exec())

if __name__ == '__main__':
    main()
